#! /usr/bin/python3

from operator import itemgetter
import collections
import decimal
from decimal import Decimal
from pprint import pprint

decimal.getcontext().prec = 4
decimal.getcontext().rounding = decimal.ROUND_HALF_UP


def _constant_factory(value):
    "The constant value factory"
    return lambda : value

def _train(corpus, n_iters = 1000):
    '''
    EM step of IBM model1 
    '''
    tgt_keys = set(
        [tgt_word for (src_sent, tgt_sent) in corpus for tgt_word in tgt_sent]
    )

    # the translation
    t = collections.defaultdict(_constant_factory(Decimal(1/len(tgt_keys))))

    for it in range(n_iters):
        cnt = collections.defaultdict(Decimal)
        tot = collections.defaultdict(Decimal)
        src_tot = collections.defaultdict(Decimal)

        for (src_sent, tgt_sent) in corpus:
            # normalization
            for src_word in src_sent:
                src_tot[src_word] = Decimal()
                for tgt_word in tgt_sent:
                    src_tot[src_word] += t[(src_word, tgt_word)]
            
            for src_word in src_sent:
                for tgt_word in tgt_sent:
                    cnt[(src_word, tgt_word)] += t[(src_word, tgt_word)] / src_tot[src_word]
                    tot[tgt_word] += t[(src_word, tgt_word)] / src_tot[src_word]
        
        for (src_word, tgt_word) in cnt.keys():
            t[(src_word, tgt_word)] = cnt[(src_word, tgt_word)] / tot[tgt_word]
    
    return t

def train(pair_sents, n_iters = 1000):
    split_pair_sents = [(src.split(), tgt.split()) for (src, tgt) in pair_sents]
    return _train(split_pair_sents, n_iters)

def test_train_loop1():
    pair_sents = [("the house", "das Haus"),
                  ("the book", "das Buch"),
                  ("a book", "ein Buch"),
                  ]
    #t0 = train(sent_pairs, loop_count=0)
    t1 = train(pair_sents, 2)

    loop1 = [(('house', 'Haus'), Decimal("0.5")),
             (('book', 'ein'), Decimal("0.5")),
             (('the', 'das'), Decimal("0.5")),
             (('the', 'Buch'), Decimal("0.25")),
             (('book', 'Buch'), Decimal("0.5")),
             (('a', 'ein'), Decimal("0.5")),
             (('book', 'das'), Decimal("0.25")),
             (('the', 'Haus'), Decimal("0.5")),
             (('house', 'das'), Decimal("0.25")),
             (('a', 'Buch'), Decimal("0.25"))]
    # assertion
    # next assertion doesn't make sence because
    # initialized by defaultdict
    #self.assertEqual(self._format(t0.items()), self._format(loop0))
    pprint(t1)
    # assert set(t1.items()) == set(loop1)

def show_train_loop3():
    pair_sents = [("the house", "das Haus"),
                  ("the book", "das Buch"),
                  ("a book", "ein Buch"),
                  ]

    for i in range(4):
        t = train(pair_sents, i)
        print('iter',i)
        pprint(t)
if __name__ == '__main__':
    # test_train_loop1()
    show_train_loop3()
